import numpy as np
from PIL import Image
import torch
import torch.nn as nn
import torch.nn.functional as F
import gradio as gr


# 定义与训练时相同的 CNN 模型架构
class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(64 * 7 * 7, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 64 * 7 * 7)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# 假设模型已经定义
model = CNN()  # 请确保使用与训练时相同的模型定义
# 加载模型并确保在 CPU 上加载，同时设置 weights_only=True
model.load_state_dict(torch.load('models/cnn_mnist_model.pth', map_location=torch.device('cpu'), weights_only=True))
model.eval()  # 设置为评估模式


# 定义图像标准化参数
image_mean = 0.5  # MNIST 数据集的平均值
image_std = 0.5   # MNIST 数据集的标准差

# 图像预处理函数
def preprocess_image(input_data):
    if isinstance(input_data, dict):
        image_data = input_data.get('composite', None)
        if image_data is None:
            image_data = input_data.get('image', None)
    else:
        image_data = input_data

    if image_data is None or np.all(image_data == 0):
        print("Empty or invalid image data received.")
        return None

    # 将图像转换为灰度图像并调整大小
    pil_image = Image.fromarray(image_data).convert('L')
    pil_image = pil_image.resize((28, 28), Image.LANCZOS)

    # 将图像数据转换为numpy数组
    image_array = np.array(pil_image).astype(float)  # 确保转换为浮点数

    # 标准化
    image_array = (image_array / 255.0 - image_mean) / image_std

    # 将图像数据转换为适合 CNN 的格式 (1, 1, 28, 28)
    image_array = image_array.reshape(1, 1, 28, 28)  # batch_size=1, channels=1, height=28, width=28

    return image_array, np.array(pil_image)  # 返回处理后的图像


# 定义预测函数
def predict_digit(image):
    if image is None:
        return {}, "", None

    # 预处理图像，并获取 28x28 图像
    processed_image, processed_image_28x28 = preprocess_image(image)

    if processed_image is None:
        return {}, "Invalid image", None

    # 将数据移到设备上（如果有可用的 GPU）
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    processed_image = torch.tensor(processed_image, dtype=torch.float32).to(device)  # 设置数据类型为 float32

    # 进行预测
    with torch.no_grad():
        output = model(processed_image)
        probabilities = nn.functional.softmax(output, dim=1).cpu().numpy()  # 获取概率
        prediction = output.argmax(dim=1).cpu().numpy()[0]  # 获取预测类别

    # 创建结果字典
    results = {str(i): float(prob) for i, prob in enumerate(probabilities[0])}

    # 将预处理后的图像转换为可以显示的格式
    pil_image = Image.fromarray(processed_image_28x28)

    return results, str(prediction), pil_image


# 创建 Gradio 接口
iface = gr.Interface(
    fn=predict_digit,
    inputs=gr.Sketchpad(
        brush=gr.Brush(colors=["#FFFFFF"], color_mode="fixed"),  # 画笔颜色设为白色
        label="在黑色画布上绘制数字",
        height=280,
        width=280
    ),
    outputs=[
        gr.Label(num_top_classes=3, label="预测概率"),
        gr.Textbox(label="预测结果"),
        gr.Image(label="预处理后的28x28图像")  # 显示预处理后的图像
    ],
    title="手写数字识别",
    description="在黑色画布上绘制白色数字，然后点击提交按钮。"
)

iface.launch()
